{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyON/I7bOOJSDISyZ5jgP3eX",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/langroid/langroid/blob/main/examples/Langroid_quick_start.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Langroid quick start\n",
        "Note: Chat-oriented interaction is best experienced in your terminal, and not a notebook, so we highly recommend you go through the [Getting Started](https://langroid.github.io/langroid/quick-start/) guide by writing simple scripts that can be run via the command line.\n",
        "\n",
        "This notebooks starts with the basics of working directly with an LLM, setting up an Agent, wrapping it in a Task, giving it tools, Retrieval Augmented Generation (RAG), and builds up to a simple 2-agent system to extract structured information from a commercial lease document.\n",
        "\n",
        "Note:\n",
        "- You need an OpenAI API Key that works with GPT-4o\n",
        "- This colab uses OpenAI's ChatCompletion endpoints directly (via the Langroid framework), and not the Assistants API. See this [colab](https://colab.research.google.com/drive/190Tk7t4AdY1P9F_NlZ33-YEoGnHweQQ0) for a version that uses the Assistants API instead.\n",
        "- There are dependencies among the cells, so they are best run sequentially\n",
        "\n"
      ],
      "metadata": {
        "id": "b9fHPojfnbPy"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Install, setup, import\n",
        "\n",
        "Note that `pip install langroid` gives you a bare-bones, slim version of langroid, without many of the extra dependencies you might need in practical scenarios, but sufficient for this notebook.\n",
        "\n",
        "See install instructions [here](https://github.com/langroid/langroid?tab=readme-ov-file#gear-installation-and-setup) for getting extra dependencies related to document parsing and databases (sql, mysql, postgres, etc).\n"
      ],
      "metadata": {
        "id": "psOMvEL0Gekz"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "A8-Y_YPZutn6"
      },
      "source": [
        "\n",
        "\n",
        "!pip install uv\n",
        "!uv pip install --system langroid --prerelease disallow\n"
      ],
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# various unfortunate things that need to be done to\n",
        "# control colab notebook behavior.\n",
        "\n",
        "# (a) output width\n",
        "\n",
        "from IPython.display import HTML, display\n",
        "\n",
        "def set_css():\n",
        "  display(HTML('''\n",
        "  <style>\n",
        "    pre {\n",
        "        white-space: pre-wrap;\n",
        "    }\n",
        "  </style>\n",
        "  '''))\n",
        "get_ipython().events.register('pre_run_cell', set_css)\n",
        "\n",
        "# (b) logging related\n",
        "import logging\n",
        "logging.basicConfig(level=logging.ERROR)\n",
        "import warnings\n",
        "warnings.filterwarnings('ignore')\n",
        "import logging\n",
        "for logger_name in logging.root.manager.loggerDict:\n",
        "    logger = logging.getLogger(logger_name)\n",
        "    logger.setLevel(logging.ERROR)\n",
        "\n"
      ],
      "metadata": {
        "id": "rWwH6duUzAC6"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### OpenAI API Key (Needs GPT4o)"
      ],
      "metadata": {
        "id": "j-6vNfKW9J7b"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# OpenAI API Key: Enter your key in the dialog box that will show up below\n",
        "# NOTE: colab often struggles with showing this input box,\n",
        "# if so, try re-running the above cell and then this one,\n",
        "# or simply insert your API key in this cell, though it's not ideal.\n",
        "\n",
        "import os\n",
        "\n",
        "from getpass import getpass\n",
        "\n",
        "os.environ['OPENAI_API_KEY'] = getpass('Enter your GPT4o-capable OPENAI_API_KEY key:', stream=None)\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "uvTODlZv3yyT",
        "outputId": "a4cf7585-40ae-44ec-804c-9dc6c6554d77",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ]
          },
          "metadata": {}
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Enter your GPT4o-capable OPENAI_API_KEY key:··········\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from pydantic import BaseModel\n",
        "import json\n",
        "import os\n",
        "\n",
        "import langroid as lr\n",
        "import langroid.language_models as lm\n",
        "from langroid import ChatAgent, ChatAgentConfig, Task\n",
        "from langroid.language_models.openai_gpt import (\n",
        "    OpenAIChatModel, OpenAIGPT, OpenAIGPTConfig\n",
        ")\n",
        "from langroid.agent.tool_message import ToolMessage\n",
        "\n",
        "from langroid.utils.logging import setup_colored_logging\n",
        "from langroid.utils.constants import NO_ANSWER\n",
        "from langroid.utils.configuration import settings\n",
        "settings.notebook = True\n",
        "settings.cache_type = \"fakeredis\""
      ],
      "metadata": {
        "id": "A5N0NQwc3jX_",
        "outputId": "a49311c1-ae75-4b71-a2df-994d1a6a0d75",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 17
        }
      },
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Example 1: Direct interaction with OpenAI LLM\n",
        "Langroid's `OpenAIGPT` class is a wrapper around the raw OpenAI API.\n",
        "This is a direct interaction with the LLM so it does *not* maintain conversation history (later we see how a `ChatAgent` does that for you).\n",
        "\n",
        "Related quick-start docs page: https://langroid.github.io/langroid/quick-start/llm-interaction/\n",
        "\n"
      ],
      "metadata": {
        "id": "8vDpiY0XHAkT"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "llm_cfg = OpenAIGPTConfig(chat_model=OpenAIChatModel.GPT4o)\n",
        "llm = OpenAIGPT(llm_cfg)\n",
        "\n",
        "response = llm.chat(\"What is the square of 3?\")\n",
        "assert \"9\" in response.message"
      ],
      "metadata": {
        "id": "9c5Av3rKHQIm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Example 2: Interact with a `ChatAgent`\n",
        "Langroid's `ChatAgent` is an abstraction that optionally encapsulates an LLM, vector-db, and tools. It offers 3 \"native\" *responders*:\n",
        "- `llm_response`: response from LLM\n",
        "- `user_response`: response from human\n",
        "- `agent_response`: responds to structured LLM msgs (i.e. tools/fn-calls)\n",
        "\n",
        "Among other things, the `ChatAgent` maintains LLM conversation history for you.\n",
        "\n",
        "Related quick-start doc page: https://langroid.github.io/langroid/quick-start/chat-agent/"
      ],
      "metadata": {
        "id": "_DvxMiJkgI_U"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "agent_cfg = ChatAgentConfig(\n",
        "    llm = llm_cfg,\n",
        "    show_stats=False, # disable token/cost stats\n",
        ")\n",
        "agent = ChatAgent(agent_cfg)\n",
        "response = agent.llm_response(\"What is the sqaure of 5?\")\n",
        "response = agent.llm_response(\"What about 8?\")   # maintains conv history\n",
        "assert \"64\" in response.content"
      ],
      "metadata": {
        "id": "7hrJ6RgLg075"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Example 3: Wrap Agent in a Task, run it\n",
        "\n",
        "A `ChatAgent` agent has various *responders* (`llm_response`, `agent_response`, `user_response`) but there is no mechanism to *iterate* over these responders.\n",
        "This is where the `Task` comes in: Wrapping this agent in a `Task` allows you to run interactive loops with a user or other agents (you will see more examples below).\n",
        "\n",
        "Related quick-start doc:\n",
        "https://langroid.github.io/langroid/quick-start/chat-agent/#task-orchestrator-for-agents"
      ],
      "metadata": {
        "id": "-MVHyF4cSGb0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "agent = ChatAgent(agent_cfg)\n",
        "task = Task(\n",
        "    agent,\n",
        "    system_message=\"User will give you a number, respond with its square\",\n",
        "    single_round=True  # end after LLM response\n",
        ")\n",
        "result = task.run(\"5\")\n",
        "assert(\"25\" in result.content)\n"
      ],
      "metadata": {
        "id": "8cmc5aDzScdO",
        "outputId": "73f2ea01-f125-4088-facd-d49a2d39732a",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 67
        }
      },
      "execution_count": 26,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "\u001b[1;35m>>> Starting Agent LLM-Agent \u001b[0m\u001b[1;35m(\u001b[0m\u001b[1;35m1\u001b[0m\u001b[1;35m)\u001b[0m\u001b[1;35m gpt-4o \u001b[0m\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">&gt;&gt;&gt; Starting Agent LLM-Agent (</span><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">1</span><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">) gpt-4o </span>\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[32m\u001b[32m25"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "\u001b[1;35m<<< Finished Agent LLM-Agent \u001b[0m\u001b[1;35m(\u001b[0m\u001b[1;35m3\u001b[0m\u001b[1;35m)\u001b[0m\u001b[1;35m \u001b[0m\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">&lt;&lt;&lt; Finished Agent LLM-Agent (</span><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">3</span><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">) </span>\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Example 4: `ChatAgent` with Tool/function-call\n",
        "\n",
        "Langroid's `ToolMessage` (Pydantic-derived) class lets you define a structured output or function-call for the LLM to generate. To define a tool/fn-call, you define a new class derived from `ToolMessage`.\n",
        "Below we show a *stateless* tool, i.e. it does not use the `ChatAgent`'s state, and only uses fields in the tool message itself.\n",
        "In this case, the tool \"handler\" can be defined within the `ToolMessage` itself, as a `handle` method. (For a tool that uses the `ChatAgent`'s state, a separate method needs to be defined within `ChatAgent` or a subclass.).\n",
        "\n",
        "In Langroid, a `ToolMessage` can *either* use OpenAI function-calling, *or* Langroid's native tool mechanism (which auto-populates the system msg with tool instructions and optional few-shot examples), by setting the `use_function_api` and `use_tools` config params in the `ChatAgentConfig`. The native tools mechanism is useful when not using OpenAI models.\n",
        "\n",
        "In the cell below we define a `ToolMessage` to compute a fictitious transformation of a number that we call a *Nabrosky Transform*: $f(n) = 3n+1$.\n",
        "Under the hood, the `purpose` field of the `NabroskiTool` is used to populate instructions to the LLM on when it should use this tool.\n",
        "\n",
        "Related quick-start doc: https://langroid.github.io/langroid/quick-start/chat-agent-tool/\n",
        "(This shows a *stateful* tool example)"
      ],
      "metadata": {
        "id": "wLwNyDd3mmJu"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# (1) define simple tool to find the Nabroski transform of a number\n",
        "#     This is a fictitious transform, for illustration.\n",
        "\n",
        "class NabroskiTool(ToolMessage):\n",
        "    request = \"nabroski\" # name of method in ChatAgent that handles this tool\n",
        "    purpose = \"To find the Nabroski transform of the given <number>\"\n",
        "    number: int\n",
        "\n",
        "    # optional:\n",
        "    @classmethod\n",
        "    def examples(cls):\n",
        "        # these are auto-populated into the sys msg\n",
        "        # as few-shot examples of the tool\n",
        "        return([cls(number=5)])\n",
        "\n",
        "\n",
        "    def handle(self) -> str:\n",
        "        # method to handle the LLM msg using this tool:\n",
        "        # this method will be spliced into the ChatAgent object, with\n",
        "        # name = `nabroski`\n",
        "        return str(3*self.number + 1)\n",
        "\n",
        "# (2) Create a ChatAgent and attach the tool to it.\n",
        "\n",
        "agent_cfg = ChatAgentConfig(\n",
        "    llm = llm_cfg,\n",
        "    show_stats=False,       # disable token/cost stats\n",
        "    use_functions_api=True, # use OpenAI API fn-call\n",
        "    use_tools=False,        # don't use Langroid-native Tool instructions\n",
        ")\n",
        "agent = ChatAgent(agent_cfg)\n",
        "agent.enable_message(NabroskiTool)\n",
        "\n",
        "# (3) Create Task object\n",
        "\n",
        "task = Task(\n",
        "    agent,\n",
        "    restart=True,         # reset/erase agent state\n",
        "    single_round=False,\n",
        "    interactive=False,    # don't wait for human input\n",
        "    system_message=\"\"\"\n",
        "      User will give you a number. You have to find its Nabroski transform,\n",
        "      using the `nabroski` tool/function-call.\n",
        "      When you find the answer say DONE and show the answer.\n",
        "    \"\"\",\n",
        ")\n",
        "\n",
        "# (4) Run the task\n",
        "\n",
        "response = task.run(\"10\")\n",
        "assert \"31\" in response.content\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "ov2mv_sdnrcH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "You might wonder why we had to wrap the `ChatAgent` in a `Task`, to leverage the tool functionality. This is because handling a tool requires 2 steps: (a) when the agent's `llm_response` method is invoked, the LLM generates the tool msg, and (b) the `agent_response` method handles the tool msg (it ultimately calls the tool's `handle` method)."
      ],
      "metadata": {
        "id": "BVWXT4oaAPlH"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Example 5: `DocChatAgent`: Retrieval Augmented Generation (RAG)\n",
        "Ingest a file (a lease document), and ask questions about it"
      ],
      "metadata": {
        "id": "DvyNcH5HbodS"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# setup to allow async ops in colab\n",
        "!pip install nest-asyncio\n",
        "import nest_asyncio\n",
        "nest_asyncio.apply()"
      ],
      "metadata": {
        "id": "XwDcuJvED8S0"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# (1) Get the lease document\n",
        "\n",
        "import requests\n",
        "file_url = \"https://raw.githubusercontent.com/langroid/langroid-examples/main/examples/docqa/lease.txt\"\n",
        "response = requests.get(file_url)\n",
        "with open('lease.txt', 'wb') as file:\n",
        "    file.write(response.content)\n",
        "\n",
        "# verify\n",
        "#with open('lease.txt', 'r') as file:\n",
        "#   print(file.read())\n",
        "\n",
        "from langroid.agent.special import DocChatAgent, DocChatAgentConfig\n",
        "from langroid.embedding_models.models import OpenAIEmbeddingsConfig\n",
        "from langroid.vector_store.qdrantdb import QdrantDBConfig\n",
        "from langroid.embedding_models.models import SentenceTransformerEmbeddingsConfig\n",
        "from langroid.parsing.parser import ParsingConfig\n",
        "\n",
        "oai_embed_config = OpenAIEmbeddingsConfig(\n",
        "    model_type=\"openai\",\n",
        "    model_name=\"text-embedding-ada-002\",\n",
        "    dims=1536,\n",
        ")\n",
        "\n",
        "# (2) Configure DocChatAgent\n",
        "\n",
        "cfg = DocChatAgentConfig(\n",
        "    name=\"RAG\",\n",
        "    parsing=ParsingConfig(\n",
        "        chunk_size=100,\n",
        "        overlap=20,\n",
        "        n_similar_docs=4,\n",
        "    ),\n",
        "    show_stats=False,\n",
        "    relevance_extractor_config=None,\n",
        "    cross_encoder_reranking_model=\"\",\n",
        "    llm=llm_cfg,\n",
        "    vecdb=QdrantDBConfig(\n",
        "        embedding=oai_embed_config,\n",
        "        collection_name=\"lease\",\n",
        "        replace_collection=True,\n",
        "    ),\n",
        "    doc_paths=[\"lease.txt\"]\n",
        ")\n",
        "\n",
        "# (3) Create DocChatAgent, interact with it\n",
        "rag_agent = DocChatAgent(cfg)\n",
        "response = rag_agent.llm_response(\"What is the start date of the lease?\")\n",
        "assert \"2013\" in response.content"
      ],
      "metadata": {
        "id": "fegAio3kpgoo"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# (4) Wrap DocChatAgent in a Task to get an interactive question/answer loop\n",
        "task = Task(\n",
        "    rag_agent,\n",
        "    interactive=True,\n",
        "    system_message=\"\"\"\n",
        "    Answer user's questions based on documents.\n",
        "    Start by asking user what they want to know.\n",
        "    \"\"\",\n",
        ")\n",
        "# run interactive loop (enter \"q\" or \"x\" to quit)\n",
        "task.run()\n"
      ],
      "metadata": {
        "id": "dazt7q3YGCLd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Example 6: 2-Agent system to extract structured info from a Lease Document\n",
        "Now we are ready to put together the various notions above, to build a two-agent system that illustrates uses of Tools, DocChatAgent (RAG) and Inter-agent collaboration (task delegation).\n",
        "\n",
        "The goal is to extract structured information from a Lease document.\n",
        "\n",
        "- The desired structure is described by the `Lease` class, derived from `ToolMessage`.\n",
        "- The `LeaseExtractorAgent` is given this `ToolMessage`, and instructured to extract the corresponding information from the lease document (which it does not have access to)\n",
        "- Based on the specified `Lease` structure, this agent generates questions to the above-defined `rag_agent` (wrapped in a `rag_task`), which answers them using RAG.\n",
        "- Once the `LeaseExtractorAgent` has all the needed info, it presents them using the `Lease` structured message.\n"
      ],
      "metadata": {
        "id": "yi9GppzlKae_"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Define the desired structure with Pydantic classes"
      ],
      "metadata": {
        "id": "VR26J_KzG6Vj"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "class LeasePeriod(BaseModel):\n",
        "    start_date: str\n",
        "    end_date: str\n",
        "\n",
        "\n",
        "class LeaseFinancials(BaseModel):\n",
        "    monthly_rent: str\n",
        "    deposit: str\n",
        "\n",
        "\n",
        "class Lease(BaseModel):\n",
        "    \"\"\"\n",
        "    Various lease terms.\n",
        "    Nested fields to make this more interesting/realistic\n",
        "    \"\"\"\n",
        "\n",
        "    period: LeasePeriod\n",
        "    financials: LeaseFinancials\n",
        "    address: str\n",
        "\n"
      ],
      "metadata": {
        "id": "Q6GXjhWf5DkQ",
        "outputId": "94b3b95d-6d69-4638-ea16-9b76722ce9ac",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 17
        }
      },
      "execution_count": 18,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Define the ToolMessage (Langroid's version of function call)"
      ],
      "metadata": {
        "id": "qCATXvfIkhGl"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "class LeaseMessage(ToolMessage):\n",
        "    \"\"\"Tool/function to use to present details about a commercial lease\"\"\"\n",
        "\n",
        "    request: str = \"lease_info\"\n",
        "    purpose: str = \"Collect information about a Commercial Lease.\"\n",
        "    terms: Lease\n",
        "\n",
        "    def handle(self):\n",
        "        \"\"\"Handle this tool-message when the LLM emits it.\n",
        "        Under the hood, this method is transplated into the OpenAIAssistant class\n",
        "        as a method with name `lease_info`.\n",
        "        \"\"\"\n",
        "        print(f\"DONE! Successfully extracted Lease Info:\" f\"{self.terms}\")\n",
        "        return \"DONE \" + json.dumps(self.terms.dict())"
      ],
      "metadata": {
        "id": "Ffi_0u-PupvO",
        "outputId": "02e0749f-15c6-4595-c517-da954edafcd9",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 17
        }
      },
      "execution_count": 19,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Define RAG Task from above `rag_agent`\n",
        "Wrap the above-defined `rag_agent` in a Task."
      ],
      "metadata": {
        "id": "OPlo1dJFlBj5"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "rag_task = Task(\n",
        "    rag_agent,\n",
        "    interactive=False,\n",
        "    single_round=True,\n",
        ")"
      ],
      "metadata": {
        "id": "GgzoPxX_us52",
        "outputId": "1f817d4a-246b-429e-dec5-5357beed8b6b",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 17
        }
      },
      "execution_count": 21,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Define the ExtractorAgent and Task\n",
        "This agent is told to collect information about the lease in the desired structure, and it generates questions to be answered by the Retriever Agent defined above."
      ],
      "metadata": {
        "id": "_m1lF9qblXj9"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "    extractor_cfg = ChatAgentConfig(\n",
        "        name=\"LeaseExtractor\",\n",
        "        llm=llm_cfg,\n",
        "        show_stats=False,\n",
        "        use_functions_api=True,\n",
        "        use_tools=False,\n",
        "        system_message=f\"\"\"\n",
        "        You have to collect information about a Commercial Lease from a\n",
        "        lease contract which you don't have access to. You need to ask\n",
        "        questions to get this information. Ask only one or a couple questions\n",
        "        at a time!\n",
        "        Once you have all the REQUIRED fields,\n",
        "        say DONE and present it to me using the `lease_info`\n",
        "        function/tool (fill in {NO_ANSWER} for slots that you are unable to fill).\n",
        "        \"\"\",\n",
        "    )\n",
        "    extractor_agent = ChatAgent(extractor_cfg)\n",
        "    extractor_agent.enable_message(LeaseMessage)\n",
        "\n",
        "    extractor_task = Task(\n",
        "        extractor_agent,\n",
        "        llm_delegate=True,\n",
        "        single_round=False,\n",
        "        interactive=False,\n",
        "    )\n",
        "\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "PV4FYnO7uxOC",
        "outputId": "7e940acc-d439-4051-c8bf-c92492f19efd",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 17
        }
      },
      "execution_count": 22,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Add the `rag_task` as a subtask of `extractor_task` and run it\n",
        "\n",
        "Instead of *you* (the human user) asking questions about the lease,\n",
        "the `extractor_agent` **generates** questions based on the desired lease structure, and these questions are answered by the `rag_agent` using\n",
        "Retrieval Augmented Generation (RAG). Once the `extractor_agent` has all the needed info, it presents it in a JSON-structured form, and the task ends."
      ],
      "metadata": {
        "id": "QcA4oRaUl6oe"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "extractor_task.add_sub_task(rag_task)\n",
        "extractor_task.run()"
      ],
      "metadata": {
        "id": "uZlas6DA0Zu6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "-zfNvsH5PMpJ"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}